import os, sys

import OpenEXR
import Imath
import math
import time

import numpy
from numpy import array
np = numpy
import myextension
    
def readExr(exrfile): 

    

    exrfile = str(exrfile)

    file = OpenEXR.InputFile(exrfile)

    pt = Imath.PixelType(Imath.PixelType.FLOAT)

    dw = file.header()['dataWindow']

    size = (dw.max.x - dw.min.x + 1, dw.max.y - dw.min.y + 1)

    R,G,B = [numpy.fromstring( file.channel(c, pt), dtype = numpy.float32) for c in ("R","G","B")]

    R = R.reshape(size[1],size[0]).T

    G = G.reshape(size[1],size[0]).T

    B = B.reshape(size[1],size[0]).T

    L = 0.2125*R + 0.7154*G + 0.0721*B

    return R,G,B,L,size



def createNewOutputImage(fileName,outputR,outputG,outputB, sz):

    print fileName

    (strR, strG, strB) = [Channel.tostring() for Channel in (outputR, outputG, outputB)]

    patchOut  = OpenEXR.OutputFile(fileName, OpenEXR.Header(sz[0], sz[1]))

    patchOut.writePixels({'R' : strR, 'G' : strG, 'B' : strB})


def positionFeature(labels, Nth):
    
    tmp = np.where(labels == Nth)
    
    return  np.mean(tmp[0]),np.std(tmp[0]),np.mean(tmp[1]),np.std(tmp[1])    

def colorFeature(L,labels, Nth):
    
    tmp = np.where(labels == Nth, L)
    return np.mean(tmp),np.std(tmp)

def features( L,labels, Nth ):
    
    tmp = np.where(labels == Nth)

    x_pos = tmp[0]
    y_pos = tmp[1]
    col   = L[np.where(labels == Nth)] 

    return np.cov(np.vstack((x_pos,y_pos,col))),np.mean(tmp[0]),np.mean(tmp[1]),np.mean(col)

def featuresRGB( R, G, B, labels, Nth ):
    
    tmp = np.where(labels == Nth)
    
    x_pos = np.array(tmp[0],dtype=np.double)
  
    y_pos = np.array(tmp[1],dtype=np.double)

    _R = R[tmp]
    _G = G[tmp]
    _B = B[tmp]
    
    return np.cov(np.vstack((x_pos,y_pos,_R,_G,_B))),np.mean(x_pos,dtype=np.double),np.mean(y_pos,dtype=np.double),np.mean(_R,dtype=np.double),np.mean(_G,dtype=np.double),np.mean(_B,dtype=np.double)

def proba_dl_thetam(dl, miu, covarMatrix):
    
    V = np.matrix(dl - np.matrix(miu))
    exp = V* np.matrix(np.linalg.inv(covarMatrix))*V.T
    
    return np.exp(-0.5*exp)/(np.power(2*np.pi,5)*np.linalg.det(covarMatrix)) 


def prob_m_dl(m,pi_m, dl, miuArray,covarMatrixArray):

    tmp1 = pi_m * proba_dl_thetam(dl,miuArray[m],covarMatrixArray[m])
    
    if m == 0:
        tmp2 = pi_m * proba_dl_thetam(dl,miuArray[0],covarMatrixArray[0]) + (1-pi_m) * proba_dl_thetam(dl,miuArray[1],covarMatrixArray[1])
    elif m == 1:
        tmp2 = (1-pi_m) * proba_dl_thetam(dl,miuArray[0],covarMatrixArray[0]) + pi_m * proba_dl_thetam(dl,miuArray[1],covarMatrixArray[1])

    return float(tmp1/tmp2)

def pi_m__t_plus_1(R,G,B,m,pi_m,miuArray,covarMatrixArray):
    
    pixels_count = R.shape[0]*R.shape[1]
    sum = 0
    for x in xrange(R.shape[0]):
        for y in xrange(R.shape[1]):
        
            dl = np.matrix([x,y,R[x][y],G[x][y],B[x][y]])
            sum = sum + prob_m_dl(m,pi_m[m], dl, miuArray,covarMatrixArray)
    
    return sum/pixels_count

def miu_m__t_plus_1(R,G,B,m,pi_m,miuArray,covarMatrixArray):
    
    sum1 = np.matrix([0,0,0,0,0])
    sum2 =0
    for x in xrange(R.shape[0]):
        for y in xrange(R.shape[1]):
            
            dl = np.matrix([x,y,R[x][y],G[x][y],B[x][y]])
            tmp = prob_m_dl(m,pi_m, dl, miuArray,covarMatrixArray)

            sum1 = sum1 + dl * tmp
            sum2 = sum2 + tmp
    
    return sum1/sum2


def covarMatrix_m__t_plus_1(R,G,B,m,pi_m,miuArray,covarMatrix):
    
    sum1 = np.matrix(np.zeros((5,5)))
    sum2 =0
    for x in xrange(R.shape[0]):
        for y in xrange(R.shape[1]):
            
            dl = np.matrix([x,y,R[x][y],G[x][y],B[x][y]])
            miu_m = np.matrix(miuArray)
            tmp = prob_m_dl(m,pi_m, dl, miuArray,covarMatrix)
            sum1 = sum1 + (dl - miu_m).T * (dl - miu_m) * tmp 
            sum2 = sum2 + tmp
    
    return sum1/sum2


if __name__ == "__main__":
    
    if len(sys.argv)  < 2:
	print "no image input"
        sys.exit(1)
        
    image = sys.argv[1]
    n_labels = 2     
    R,G,B,L,size = readExr(image)
    ########
    R =  np.array(R,dtype = np.double)
    G =  np.array(G,dtype = np.double)
    B =  np.array(B,dtype = np.double)
    L =  np.array(L,dtype = np.double)
    ########
    
    print image,size
    #initialisation of labels
    #labels = np.array(np.random.randint(n_labels,size=size),dtype=np.double)
    
    
    labels = np.ones(size,dtype=np.double)
    # sunflower
    labels[115:293,(492-378):(492-327)] = 0
    labels[156:264,(492-182):(492-128)] = 0
    labels[78:135,(492-302):(492-201)] = 0
    labels[287:324,(492-322):(492-207)] = 0
    
    
    #eye.exr
    #labels[81:142,(185-103):(185-49)] = 0
    
    #eye_small.exr
    #labels[15:29,(36-20):(36-9)] = 0
    
    #Pixar05.exr
    #labels[119:205,(702-227):(702-63)] = 0
    #labels[446:495,(702-438):(702-420)] = 0
    
    #pixar.exr
    #labels[50:91,(146-92):(146-44)] = 0
    
    #pixar_creation.exr
    #labels[552:615,(511-229):(511-190)] = 0
    
    #labels[62:97,(43-39):(43-29)] = 0
    #labels[23:39,(59-42):(59-12)] = 0
    
    
    createNewOutputImage("../../images/label0.exr",np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T, size)
    
    maxflow = 0
    for k in xrange(0): 
        
        inversedCovarianceMatrixArray = []
        miuArray = []
        lnCovarMatDet = []
        try:
            for i in xrange(n_labels):
                
                covarMatrix, x,y,lum = features(L,labels,i)
                inversedCovarianceMatrixArray.append(np.linalg.inv(covarMatrix))
                miuArray.append((x,y,lum))    
                lnCovarMatDet.append(np.log(np.sqrt(8 * np.pi*np.pi*np.pi * np.linalg.det(covarMatrix))))
            
            inversedCovarianceMatrixArray   = np.array(inversedCovarianceMatrixArray).reshape((n_labels,3,3))
            miuArray                        = np.array(miuArray).reshape((n_labels,3))
            lnCovarMatDet                   = np.array(lnCovarMatDet).reshape(n_labels)
        except:
            print "exception"
            break
        
        flow = myextension.quickGraphCut(n_labels, L, labels, miuArray, inversedCovarianceMatrixArray,lnCovarMatDet)
        createNewOutputImage("../../images/label"+str(k+1)+".exr",np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T, size)
    
        if flow > maxflow:
            maxflow = flow
        else:
            pass
            #break   
        
    
    
    for k in xrange(6): 
        
        inversedCovarianceMatrixArray = []
        miuArray = []
        lnCovarMatDet = []
        covarMatrixArray = []
  
        for i in xrange(n_labels):
                
            covarMatrix, x, y, r, g, b = featuresRGB(R,G,B,labels,i)
            
            inversedCovarianceMatrixArray.append(np.linalg.inv(covarMatrix))
                
            miuArray.append((x,y,r,g,b))  
            
            lnCovarMatDet.append(np.log(np.sqrt( np.power(2*np.pi,5)  * np.linalg.det(covarMatrix))))
                
            covarMatrixArray.append(covarMatrix)
            
        inversedCovarianceMatrixArray   = np.array(inversedCovarianceMatrixArray,dtype = np.double).reshape((n_labels,5,5))                
        miuArray                        = np.array(miuArray,dtype = np.double).reshape((n_labels,5))
        lnCovarMatDet                   = np.array(lnCovarMatDet,dtype = np.double).reshape(n_labels)

        flow = myextension.quickGraphCut(n_labels, R,G,B, labels, miuArray, inversedCovarianceMatrixArray,lnCovarMatDet)
        if flow > maxflow:
                maxflow = flow
        else:
            pass
            #break
        createNewOutputImage("../../images/label"+str(k+1)+".exr",np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T, size)
    
            
    Pr_t_m = list()
    Pr_t_m.append( 0.5)
    Pr_t_m.append( 0.5) 

    Pr_t_m = np.array(Pr_t_m,dtype = np.double)
            
    Pr_t_plus_1_m = list()
    Pr_t_plus_1_m.append( 0.5)
    Pr_t_plus_1_m.append( 0.5)
            
    Pr_t_plus_1_m = np.array(Pr_t_plus_1_m,dtype = np.double)
            
    inversedCovarianceMatrixArray = []
    lnCovarMatDet = [1,1]
    miuArray = []
    covarMatrixArray = []
          
    for i in xrange(n_labels):
        
        covarMatrix, x, y, r, g, b = featuresRGB(R,G,B,labels,i)
        miuArray.append((x,y,r,g,b))    
        covarMatrixArray.append(covarMatrix)
        inversedCovarianceMatrixArray.append(np.linalg.inv(covarMatrix))

            
    miuArray = np.array(miuArray,dtype = np.double).reshape((n_labels,5))
    
    covarMatrixArray = np.array(covarMatrixArray,dtype = np.double).reshape((n_labels,5,5))

    inversedCovarianceMatrixArray = np.array(inversedCovarianceMatrixArray,dtype = np.double).reshape((n_labels,5,5))

    old_miuArray = np.array(miuArray,dtype = np.double).reshape((n_labels,5))

    new_miuArray = np.array(miuArray,dtype = np.double).reshape((n_labels,5))

    old_covarMatrixArray = np.array(covarMatrixArray,dtype = np.double).reshape((n_labels,5,5))
    
    new_covarMatrixArray = np.array(covarMatrixArray,dtype = np.double).reshape((n_labels,5,5))
        
    covarMatDet = np.array([np.linalg.det(old_covarMatrixArray[0]),np.linalg.det(old_covarMatrixArray[1])],dtype = np.double)

    for k in xrange(0): 
        
        for i in xrange(n_labels):
            pass
            """ 
            #covarMatrix, x, y, r, g, b = featuresRGB(R,G,B,labels,i)
            #miu = np.matrix([x,y,r,g,b]) 
            
            miu = np.matrix( old_miuArray[i])
            covarMatrix = old_covarMatrixArray[i]
            
            # New Miu
            sum1 = np.matrix([0,0,0,0,0])
            sum2 =0
            #New covarMatrix
            sum3 = np.matrix(np.zeros((5,5)))
            #New pr_t_m
            pixels_count = R.shape[0]*R.shape[1]
           
            
            #dl = [[x,y,R[x,y],G[x,y],B[x,y]] for x in xrange(R.shape[0]) for y in xrange(R.shape[1]) ]
            
            for x in xrange(R.shape[0]):
                for y in xrange(R.shape[1]):
            
                    dl = np.matrix([x,y,R[x,y],G[x,y],B[x,y]])
                    
                    tmp = prob_m_dl(i,Pr_t_m[i], dl, old_miuArray,old_covarMatrixArray) # reponsibility P(m | dl)
            
                    sum1 = sum1 + dl * tmp
                    sum2 = sum2 + tmp
                
                    sum3 = sum3 + (dl - miu).T * (dl - miu) * tmp 
            
            print sum2
            new_miu = sum1/sum2
            print "miu ",miu
            print "new miu",new_miu 
            miu = new_miu
            new_miuArray[i]=miu
        
            new_covarMatrix = sum3/sum2
            new_covarMatrixArray[i]=(new_covarMatrix)
            print old_covarMatrixArray[i]
            print new_covarMatrix
            Pr_t_plus_1_m[i] = sum2/pixels_count
                            
            inversedCovarianceMatrixArray[i] = np.linalg.inv(new_covarMatrix)
            lnCovarMatDet[i]=np.log(np.sqrt(np.power(2*np.pi,5) * np.linalg.det(new_covarMatrix)))
            """
    
        myextension.EMProcess(Pr_t_m, Pr_t_plus_1_m, R, G, B, old_miuArray, new_miuArray, inversedCovarianceMatrixArray, new_covarMatrixArray, covarMatDet)
        
        print old_miuArray
        print new_miuArray

        old_miuArray = np.array(new_miuArray,dtype = np.double)
        old_covarMatrixArray=np.array(new_covarMatrixArray,dtype=np.double) 
        Pr_t_m = np.array(Pr_t_plus_1_m,dtype=np.double)
        
        covarMatDet = np.array([np.linalg.det(new_covarMatrixArray[0]),np.linalg.det(new_covarMatrixArray[1])])
        
    
        inversedCovarianceMatrixArray[0] = np.linalg.inv(new_covarMatrixArray[0])
        inversedCovarianceMatrixArray[1] = np.linalg.inv(new_covarMatrixArray[1])
        
        inversedCovarianceMatrixArray   = np.array(inversedCovarianceMatrixArray,dtype = np.double).reshape((n_labels,5,5))                 
       
        miuArray                        = np.array(new_miuArray,dtype=np.double).reshape((n_labels,5))
        
        lnCovarMatDet[0] =  np.log(np.sqrt(np.power(2*np.pi,5) * np.linalg.det(old_covarMatrixArray[0])))
        lnCovarMatDet[1] =  np.log(np.sqrt(np.power(2*np.pi,5) * np.linalg.det(old_covarMatrixArray[1])))
         
        lnCovarMatDet                   = np.array(lnCovarMatDet,dtype=np.double).reshape(n_labels)

        flow = myextension.quickGraphCut(n_labels, R,G,B, labels, miuArray, inversedCovarianceMatrixArray,lnCovarMatDet)
        
        if flow > maxflow:
            maxflow = flow
        else:
            pass
            #break
        createNewOutputImage("../../images/label"+str(k+1)+".exr",np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T, size)

    
    # Here begins the alpha expansion, we extract at first the contour of the initial alpha regions as active nodes ( because inside of alpha regions, the Pott cost is 0). the initial alpha region is difined by a rectangle (labels). so the contour is four edges of the rectangle.

    alphaRegion = np.where(labels == 0) # extract alpha Region

    max_X = np.max(alphaRegion[0])
    min_X = np.min(alphaRegion[0]) 
    max_Y = np.max(alphaRegion[1])
    min_Y = np.min(alphaRegion[1]) 

    activeContour = [ [x,min_Y] for x in xrange(min_X,max_X+1) ] + [ [x,max_Y] for x in xrange(min_X,max_X+1) ] + \
                  [ [min_X,y] for y in xrange(min_Y,max_Y+1) ] + [ [min_X,y] for x in xrange(min_Y,max_Y+1) ] 
    activeContour = np.array(activeContour, dtype = np.double)

    print activeContour.shape
    
    #myextension.alphaExpansionQuickGraphCut(R,G,B,labels,miuArray,inversedCovarianceMatrixArray,lnCovarMatDet,activeContour)
    #createNewOutputImage("label_expansion.exr",np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T,np.array(labels,dtype=np.float32).T, size)






